#include "RSA_.hpp"


using namespace::std;

typedef struct RSAentry
{
	WLANAddr address;
	unsigned int publicKey;
	unsigned int privateKey;
	unsigned int nKey;
	double time;
}RSAtableEntry;

vector<RSAtableEntry> RSATable;


void handleRSA(PCryptoPacket pCryptoPacket, WLANHeader * srcHeader )
{
	char b[50];
	unsigned int c;
	unsigned int m;
	
	WLANAddr srcAddr = srcHeader->srcAddr;
        PCryptoPacket pcp = NULL;

	
	
	
	int id = checkRSATable(&srcAddr);
	
	if( id != -1)
		popFromRSATable(id);
	char * buff = pCryptoPacket->data;
	unsigned short type = pCryptoPacket->id & 0x0F;
	printf ("type=%u",type);

	// Entry [RSA_SEND], decrypt message with d
	if ( type == RSA_SEND) {
		unsigned int mes;		
		printf ("Recieved Encrypted Message from: id=%d\n",id);
		
		id = checkRSATable(&srcAddr);
		memcpy( &mes, pCryptoPacket->data, sizeof(unsigned int));

		//mes = decrypt(mes, RSATable.at(id).privateKey, RSATable.at(id).nKey);
		
		printf ("The message was: %u\n", mes);
		return;
	}

	// Entry [RSA_INIT], computer n,e,d and send n,e 
	if ( type == RSA_INIT) {
		printf ("Sending public key to: id=%d\n",id);

		// Extract Public Key
		unsigned int public_key;
		memcpy( &public_key, pCryptoPacket->data, sizeof(unsigned int)); 	
	
		// calculate shared key + add
		unsigned int timer = 1;
				
		genKeys(&srcAddr, timer);
		
		id = checkRSATable(&srcAddr);
		// generate message then send
			pcp = (CryptoPacket *)malloc(sizeof(CryptoPacket));
			if ( pcp != NULL )
			{

				pcp->id = RSA_PUBLICKEY | (RSA << 4);

				// add the data (e,n)
				memcpy( pcp->data, &RSATable.at(id).publicKey, sizeof(unsigned int) );
				memcpy( pcp->data + sizeof(unsigned int), &RSATable.at(id).nKey, sizeof(unsigned int) );

				// send it
				printf ("Sending Public Keys (n,e)\n");
				sendRSACryptoPacket( pcp, &srcAddr );


			}
			else
				printf ("PCP failed malloc");

	}

	// Entry [PUBLICKEY], store the data and send message
	else if ( type == RSA_PUBLICKEY ) { 
			printf ("Public Keys(n,e) Recieved\n");

			// Strip p, g, public_key values
			unsigned int e = 1;
			unsigned int n = 1;
			unsigned int d = 1;


			// generate own public_key
			memcpy( &e, pCryptoPacket->data, sizeof(unsigned int)); 
			memcpy( &n, pCryptoPacket->data + (sizeof(unsigned int)), sizeof(unsigned 				int)); 


			printf ("e=%u\n",e);
			printf ("n=%u\n",n); 

			RSAtableEntry newEntry;
			newEntry.address = srcAddr;
			newEntry.nKey = n;
			newEntry.privateKey = d;
			newEntry.publicKey = e;
			newEntry.time = 1;
			RSATable.push_back(newEntry);

			

			// generate message then send
			pcp = (CryptoPacket *)malloc(sizeof(CryptoPacket));
			if ( pcp != NULL )
			{

				pcp->id = RSA_PUBLICKEY | (RSA << 4);
				char temp[2];
				temp[0]= '1';
				temp[1]= '7';
				unsigned int mess = encrypt(temp, e, n);
				// add the data (this node's public key)
				memcpy( pcp->data, &mess , sizeof(unsigned int) );

				// send it
				//printf ("Sending Encrypted Data\n");
				//sendRSACryptoPacket( pcp, &srcAddr );


			}
			else
				printf ("PCP failed malloc");




		} 
	

}
	
void initiateRSA()
{
	WLANAddr destAddr [ WLAN_ADDR_LEN ];
	PCryptoPacket pcp = NULL;

	srand ( time(NULL) );	

	pcp = (CryptoPacket *)malloc(sizeof(CryptoPacket));
	if ( pcp != NULL )
	{
		for (int i = 0; i < WLAN_ADDR_LEN; i++ )
		{
			destAddr->data[i] = 0xFF;
		}

		/* If they don't already exist, establish RSA connection */
		if (checkRSATable(destAddr) == -1)
		{
			pcp->id = RSA_INIT | (RSA << 4);
			sendRSACryptoPacket( pcp, destAddr );

			printf ("Sent RSA Initialization\n");

		}
	}
}
void sendRSACryptoPacket( PCryptoPacket pCryptoPacket, WLANAddr * destAddr)
{

    char asciiAddr[32];

	char * dest = destAddr->wlan2asc();
	strcpy(asciiAddr, dest);
	printf ("Send to: %s\n", asciiAddr);
	


	void * pBuffer = malloc( sizeof(CryptoPacket) );
	if ( pBuffer != NULL )
	{
		memcpy( (char*)pBuffer, pCryptoPacket, sizeof(CryptoPacket) );
		MySendData( destAddr->data, pBuffer );
	}

	if ( pCryptoPacket != NULL )
	{
		free( pCryptoPacket );
	}

}
/*
RSA code
*/
void genKeys(WLANAddr * addr, double timer)
{
	unsigned long long p;
	unsigned long long q;
	unsigned long long n;
	unsigned int En;
	unsigned int e;
	unsigned int d;
	unsigned int sum, count=0, k;
	
	//generate 2 prime numbers
	p = genPrimeNum();
	q = genPrimeNum();
	
	//compute n = pq
	n = p*q;
	
	//compute Euler's totient function
	En = (p-1)*(q-1);
	
	//choose an int such that 1 < e < En and GCD of (e,En) = 1
	e = genCoPrime(En);
	//determine d
	d = findInverse(e,En);
	//add the new connection to table
	RSAtableEntry newEntry;
	newEntry.address = *addr;
	newEntry.nKey = n;
	newEntry.privateKey = d;
	newEntry.publicKey = e;
	newEntry.time = timer;
	RSATable.push_back(newEntry);
	
}
unsigned int encrypt(char* buff, unsigned int e, unsigned int n)
{
	unsigned int c;
	char b[50];
	unsigned int buffint = atoi(buff);
	c = modularExponentiation(buffint, e, n);
	//buff = itoa(c,b,10);
	return c;
}
unsigned int decrypt(char* buff, unsigned int d, unsigned int n)
{
	unsigned int m;
	char b[50];
	int buffint = atoi(buff);
	m = modularExponentiation(buffint, d, n);
	//buff = itoa(m,b,100);
	return m;
}
/*
RSA code end
*/
/*
Generating code Start
*/
unsigned int genCoPrime(int n)
{
	bool isCoPrime = false;
	unsigned int c;
	c = rand() % 1000 + 2;
	while(gcd(c,n) != 1)
	{
		c++;
	}
	return c;
}
int gcd(int a, int b)
{
	if(b == 0)
	{
		return a;
	}else{
		return gcd(b, a%b);
	}

}
unsigned int genPrimeNum(void)
{
	unsigned int n;
	n = rand() % 40000 + 10000;
	while(is_Prime(n) != 1)
	{
		n++;
	}
	return n;
}
int is_Prime(int n)
{
	bool prime = true;

	for ( int i = 2; (i <= n / 2) && (prime == true); i++ )
		if (n % i == 0)prime = false;

	return prime;
}


int findInverse(unsigned int u, unsigned int v)
{
	unsigned int inv, u1, u3, v1, v3, t1, t3, q;
    int iter;
    /* Step X1. Initialise */
    u1 = 1;
    u3 = u;
    v1 = 0;
    v3 = v;
    /* Remember odd/even iterations */
    iter = 1;
    /* Step X2. Loop while v3 != 0 */
    while (v3 != 0)
    {
        /* Step X3. Divide and "Subtract" */
        q = u3 / v3;
        t3 = u3 % v3;
        t1 = u1 + q * v1;
        /* Swap */
        u1 = v1; v1 = t1; u3 = v3; v3 = t3;
        iter = -iter;
    }
    /* Make sure u3 = gcd(u,v) == 1 */
    if (u3 != 1)
        return 0;   /* Error: No inverse exists */
    /* Ensure a positive result */
    if (iter < 0)
        inv = v - u1;
    else
        inv = u1;
    return inv;
}

/*
Generating code End
*/
/*Helper code:*/
int checkRSATable (WLANAddr * addr) {
	for (int i = 0; i < RSATable.size(); i++) {
		if (RSATable.at(i).address.data[0] == addr->data[0] &&
  			RSATable.at(i).address.data[1] == addr->data[1] &&		
			RSATable.at(i).address.data[2] == addr->data[2] &&		
			RSATable.at(i).address.data[3] == addr->data[3] &&
			RSATable.at(i).address.data[4] == addr->data[4] &&
			RSATable.at(i).address.data[5] == addr->data[5])
				return i;
	} 
	return -1;
}
void printRSATable()
{
	printf("Address\t\t\tPublicKey\tPrivateKey\tN Key\n");	
	for(int i =0;i<RSATable.size();i++)
	{
		printf("%s\t%u\t\t%u\t\t%u\n", RSATable.at(i).address.wlan2asc(), RSATable.at	  (i).publicKey, RSATable.at(i).privateKey, RSATable.at(i).nKey);
	}
}
void popFromRSATable (int id) {
		// Swap and pop to remove from RSA stack
		RSAtableEntry temp = {NULL,0,0,0,0};
		temp = RSATable.at(RSATable.size()-1);
		RSATable.at(RSATable.size()-1) = RSATable.at(id);
		RSATable.at(id) = temp;
		RSATable.pop_back();
}
